#!/usr/bin/python3
# -*- coding: utf8 -*-

# Copyright (c) 2021 Baidu, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Example: QuanlseScheduler
Please visit https://quanlse.baidu.com/#/doc/tutorial-scheduler for more details about this example.
"""

import math
import numpy

from Quanlse import Define
from Quanlse import BackendName, Algorithm, HardwareImplementation
from Quanlse.QuanlseEnv import QuanlseEnv
from Quanlse.QOperation.FixedGate import X, H, CZ
from Quanlse.QOperation.RotationGate import U
from Quanlse.Utils.Plot import plotScheduler, plotBarGraph
from Quanlse.Utils.Tools import computationalBasisList
from Quanlse.Utils.Operator import basis
from Quanlse.QPlatform.Utilities import dictMatrixToNumpyMatrix

# Your token:
# Please visit http://quantum-hub.baidu.com
Define.hubToken = ''

env = QuanlseEnv()
env.backend(BackendName.CloudScheduler, Algorithm.Normal, HardwareImplementation.CZ, 1.0)


# Define the qubit frequency and anharmonicity.
env.setupQReg([
    (5.805 * 2 * math.pi, -0.226 * 2 * math.pi),  # QReg 0
    (5.205 * 2 * math.pi, -0.226 * 2 * math.pi),  # QReg 1
    (4.605 * 2 * math.pi, -0.226 * 2 * math.pi),  # QReg 2
])

# Define the coupling strength, as well as the optimization bounds for the flux amplitudes
#   when the pulse acts on the corresponding qubit pairs.
env.setupCoupling([
    ((0, 1), 0.0277 * 2 * math.pi, (-3, 0)),
    ((1, 2), 0.0277 * 2 * math.pi, (-3, 0)),
    ((2, 1), 0.0277 * 2 * math.pi, (-3, 0)),
])

# Define the logic quantum circuit.
# R1
U(theta=-1.231, phi=0, lamda=0)(env.Q[0])

# X gate
X(env.Q[0])

# CNOT: 0 -> 1
# The CNOT gate is realized using CZ gate with two local single qubit gates.
H(env.Q[1])
CZ(env.Q[0], env.Q[1])
H(env.Q[1])

# X gate
X(env.Q[0])

# R2
U(theta=-0.785, phi=0, lamda=0)(env.Q[2])

# CNOT: 1 -> 2
# The CNOT gate is realized using CZ gate with two local single qubit gates.
H(env.Q[2])
CZ(env.Q[1], env.Q[2])
H(env.Q[2])

# R3
U(theta=0.785, phi=0, lamda=0)(env.Q[2])

# CNOT: 2 -> 1
# The CNOT gate is realized using CZ gate with two local single qubit gates.
H(env.Q[1])
CZ(env.Q[1], env.Q[2])
H(env.Q[1])

# Run the pulse scheduling and benchmarking.
taskResult = env.commit()
print("Benchmarking infidelity: ", taskResult["results"]["benchmark"]["infidelity"])

# Plot the pulse sequences.
plotScheduler(taskResult["results"]["scheduler"])

# Calculate the population of the eigenstates.
unitary = dictMatrixToNumpyMatrix(taskResult["results"]["benchmark"]["unitary"], valueType=complex)
finalState = (unitary @ numpy.array(basis(27, 0))).T[0]
popList = [abs(item ** 2) for item in finalState]

# Draw the population of the states in computational basis.
basis = computationalBasisList(3, 3)
plotBarGraph(basis, popList, "Population of a W state generated by QuanlseScheduler",
             "Computational Basis", "Population")
